# encoding: utf-8
"""
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    created by lane.chang on '13/05/2024'
    comment: 界面协作者
"""
import ujson
from datetime import datetime

from config import Config
from project.core.log import Colors
from project.lib.common import logger
from project.lib.common import get_redis, session_maker
from project.model.basic import User
from project.model.chatgpt import Agent as AgentModel
from project.model.chatgpt import ChatGpt


class Agent:
    """
    """
    def __init__(self, user: User,
                 user_message: str = '',
                 memories_key='interaction',
                 chat_model_name=None,
                 emb_model_name='',
                 business_name='',
                 agent_profile='',
                 user_role='',
                 is_init=False):
        """
        :param user_message: 用户话术
        """
        self.user = user
        self.user_message = user_message
        self.memories_key = memories_key
        if self.memories_key:
            self.memories_key = f'{self.memories_key}_memories_{self.user.dev_id}_{self.user.bot_id}_{self.user.user_id}'
        self.is_init = is_init
        self.chat_model_name = chat_model_name or Config.DEFAULT_MODEL
        self.emb_model_name = emb_model_name or 'text-embedding-ada-002'
        self.business_name = business_name
        self.agent_profile = agent_profile or '你是一个在酒店房间里的人工智能机器人'
        self.user_role = user_role or '客人'

    def format(self, message, _type='in'):
        """
        :param message:
        :param _type:
        :return:
        """
        now = datetime.now()

        if _type == 'in':
            ret = '对话序号:' + now.strftime('%Y%m%d%H%M%S') + f"\t\t{self.user_role}对你说:" + message
        else:
            ret = '对话序号:' + now.strftime('%Y%m%d%H%M%S') + f"\t\t你对{self.user_role}说:" + message

        return ret

    def add_user_memories(self, *args, **kwargs):
        """ 增加用户对话记录
        :param args:
        :param kwargs:
        :return:
        """
        if self.memories_key:
            rds = get_redis()
            user_memories = list()
            if rds.exists(self.memories_key):
                user_memories = rds.hget(self.memories_key, 'memories') or '[]'
                user_memories = ujson.loads(user_memories)

            for v in args:
                user_memories.append(v)

            for _, v in kwargs.items():
                user_memories.append(v)

            rds.hset(self.memories_key, 'memories', ujson.dumps(user_memories))

    def load_user_memories(self, _type='json', count=10):
        """ 加载历史对话记录
        :param _type:
        :param count: 取得最新10条记录
        :return:
        """
        rds = get_redis()
        user_memories = rds.hget(self.memories_key, 'memories') or '[]'
        user_memories = ujson.loads(user_memories)
        user_memories = user_memories[0 - count:]
        if _type == 'json':
            user_memories = ujson.dumps(user_memories, ensure_ascii=False)

        return user_memories

    def init_user_memories(self):
        """ 清空历史对话记录
        :return:
        """
        rds = get_redis()
        rds.delete(self.memories_key)

    def add_agent_io(self, input, output):
        """
        :param input:
        :param output:
        :return:
        """
        # 记录agent输入输出
        with session_maker() as session:
            AgentModel.create_modify(
                session=session,
                dev_id=self.user.dev_id,
                bot_id=self.user.bot_id,
                user_id=self.user.user_id,
                user_message=self.user_message,
                business_name=self.business_name,
                model=self.chat_model_name,
                input=input,
                output=output)

    async def async_interact(self, prompt, **kwargs):
        """
        :param prompt:
        :param kwargs:
        :return:
        """
        # 清空历史缓存
        if self.is_init:
            logger.info(f'清空历史缓存', font_color=Colors.PURPLE.value)
            self.init_user_memories()

        chat_gpt = ChatGpt(chat_model_name=self.chat_model_name, emb_model_name=self.emb_model_name)
        user_memories = kwargs.get('user_memories') or self.load_user_memories()
        kwargs.update({'user_memories': user_memories})
        prompt = prompt.format(user_message=self.user_message, agent_profile=self.agent_profile, user_role=self.user_role, **kwargs)
        response = await chat_gpt.llm_async(user_text=prompt)
        # 记录agent输入输出
        self.add_agent_io(input=prompt, output=response)
        logger.info(f'用户说: {self.user_message} Response: {response}', font_color=Colors.PURPLE.value)

        return response

    async def async_interact_streaming(self, prompt, **kwargs):
        """ 流式返回
        :param prompt:
        :param kwargs:
        :return:
        """
        # 清空历史缓存
        if self.is_init:
            logger.info(f'清空历史缓存', font_color=Colors.PURPLE.value)
            self.init_user_memories()

        chat_gpt = ChatGpt(chat_model_name=self.chat_model_name, emb_model_name=self.emb_model_name)
        user_memories = kwargs.get('user_memories') or self.load_user_memories()
        kwargs.update({'user_memories': user_memories})
        prompt = prompt.format(user_message=self.user_message, agent_profile=self.agent_profile, user_role=self.user_role, **kwargs)
        response_stream = chat_gpt.llm_async_streaming(user_text=prompt)
        # # 记录agent输入输出
        self.add_agent_io(input=prompt, output='流式')
        # logger.info(f'用户说: {self.user_message} Response: {response}', font_color=Colors.PURPLE.value)
        return response_stream


if __name__ == '__main__':

    user = User()
    agent = Agent(user, user_message='')
    # agent.add_user_memories('老人对你说: 有米饭吗?')
    # agent.add_user_memories('你对老人说: 没有，还需要其它的吗？')
    #
    # agent.add_user_memories('老人对你说: 有米饭吗?', '你对老人说: 没有，还需要其它的吗？', '老人对你说: 不需要了')

    print(agent.init_user_memories())

